from __future__ import division
from itertools import chain
import numpy as np

from aima.search import Problem
from aima.search import astar_search, best_first_graph_search

from osm_utils import Link

from our_utils import log, gamma, is_turn, get_latlon
from our_utils import timed


class WazeProblem(Problem):
    '''
    Implements the abstract class aima.search.Problem, for use in aima's
    best-first alogrithms.
    A state is a Link. an action is a Link too.
    '''
    c = 0

    def __init__(self, road_map, initial, goal):
        '@param weights may be one of "fast", "close" or a tuple of weights'
        log('{}) new problem: {}\t{}\t'.format(WazeProblem.c, initial, goal))
        self.c = WazeProblem.c
        WazeProblem.c += 1
        first = Link(initial, initial, 0, 0)
        first.distance_to_goal = road_map.JunctionDistance(initial, goal)
        super(WazeProblem, self).__init__(first, goal)
        for j in road_map.junctions:
            for link in j.links:
                link.distance_to_goal = road_map.JunctionDistance(link.target, self.goal)
        self.road_map = road_map
        log('problem', self.c, 'ready')

    def set_weights(self, weights):
        self.weights = weights
        fs = {'fast': (self._cost_fastest, self._heuristic_fastest),
              'close': (self._cost_closest, self._heuristic_closest)}
        if self.weights not in fs:
            self.f, self.h = self._cost_sigma, self._heuristic_sigma
            self.coeffs = gamma.values * self.weights
        else:
            self.f, self.h = fs[self.weights]
        return self

    def actions(self, state):
        return self.road_map.junctions[state.target].links

    def result(self, state, action):
        return action

    def goal_test(self, state):
        return state.target == self.goal

    def path_cost(self, c, state1, action, state2):
        'states and action are Links'
        'dtt is distance to target'
        return c + self.f(state1, action)

    def _cost_closest(self, from_link, to_link):
        return (to_link.distance + 1) / 1000

    def _cost_fastest(self, from_link, to_link):
        return self._cost_closest(from_link, to_link) / to_link.speed

    def _cost_simplest(self, from_link, to_link):
        rm = self.road_map
        a1 = get_latlon(rm.junctions[from_link.source])
        a2 = get_latlon(rm.junctions[from_link.target])
        a3 = get_latlon(rm.junctions[to_link.source])
        return is_turn(a2 - a1, a3 - a2)

    def _cost_sigma(self, from_link, to_link):
        costs = (self._cost_closest(from_link, to_link),
                 self._cost_fastest(from_link, to_link),
                 self._cost_simplest(from_link, to_link))
        return np.dot(self.coeffs, costs)

    def _heuristic_closest(self, node):
        return node.state.distance_to_goal

    def _heuristic_fastest(self, node):
        return node.state.distance_to_goal / self.road_map._GetMaxSpeed()

    def _heuristic_sigma(self, node):
        distance = self._heuristic_closest(node)
        max_speed = self.road_map._GetMaxSpeed()
        return np.dot(self.coeffs, (distance, distance / max_speed, 0))

    @timed
    def RunWeightedAstar(self, w):
        log('run:', self.c, ' - ', w)

        def cost(n):
            return w * n.path_cost + (1 - w) * self.h(n)
        res = best_first_graph_search(self, cost)
        res.w = w
        log('run', self.c, ' - ', w, 'finished')
        return res

    def RunAstar(self):
        return timed(astar_search)(self, self.h)


def draw_roundtrip(s, g):
    from osm_utils import CountryMap
    import plot_map as pm
    import matplotlib.pyplot as plt
    road_map = CountryMap()
    for t in ('close', 'fast', (0.4, 0.4, 0.2)):
        for ss, gg, cc in [(s, g, 'b-'), (g, s, 'r.')]:
            res = WazeProblem(road_map, ss, gg, t).RunAstar()
            pm.plot_path(road_map, (x.state for x in res.path()), cc)
        plt.show()

if __name__ == '__main__':
    draw_roundtrip(108838, 111600)
